#include "DeepAppearanceDescriptor/FeatureTensor.h"
#include "DeepAppearanceDescriptor/model.h"
#include "object_detect_alg.hpp"
#include <opencv2/opencv.hpp>
#include "KalmanFilter/tracker.h"

DETECTIONS get_detections(const object_detect_alg::bbox_t &bbox) {
    DETECTIONS detections;
    DETECTION_ROW detectionRow;
    detectionRow.tlwh = DETECTBOX(bbox.rect.x, bbox.rect.y, bbox.rect.width, bbox.rect.height);
    detectionRow.confidence = bbox.prob;   
    detections.push_back(detectionRow);
    return detections;
}

// 计算IOU（Intersection over Union）
void compare_iou(const std::vector<object_detect_alg::bbox_t> track_boxes, std::vector<object_detect_alg::bbox_t> &boxes){
    for (const auto& track_box : track_boxes) {
        for (auto& box : boxes) {            
            // 计算IOU
            cv::Rect intersection = track_box.rect & box.rect;
            float iou = static_cast<float>(intersection.area()) / (track_box.rect.area() + box.rect.area() - intersection.area());
            if (iou > 0.7) { // 如果IOU大于0.5，认为是同一个目标
                box.track_id = track_box.track_id; // 更新box为track_box              
            }
            
        }
    }
}

int main() {
    // 加载单张测试图像
    cv::Mat testImage = cv::imread("t1.jpg");
    const int nn_budget = 100;
    const float max_cosine_distance = 0.7;

    object_detect_alg *yolo=create_object_detect_alg();
    yolo->init("/home/administrator/Downloads/test_model_coco/yolov8n.onnx","yolov8.cfg","/home/administrator/Downloads/test_model_coco/obj.names");

    // 初始化特征提取器
    FeatureTensor* extractor = FeatureTensor::getInstance();
    extractor->init("./resnet50.onnx");

    tracker mytrack(max_cosine_distance, nn_budget);

    // std::cout << "OpenCV版本: " << CV_VERSION << std::endl;
    // std::cout << "编译信息: " << std::endl << cv::getBuildInformation() << std::endl;

    cv::VideoCapture cap("/home/administrator/deepsort_onnx_cpp/build/test_1.mp4");

    if (!cap.isOpened()) {
        std::cerr << "Error: Could not open video file." << std::endl;
        return -1;
    }
    // 读取视频帧
    cv::Mat frame;
    while (cap.read(frame)) {
        // 显示视频帧
        if (frame.empty()) {
            std::cerr << "Error: Could not read frame from video." << std::endl;
            break;
        }

        // 使用YOLOv8进行目标检测
        std::vector<object_detect_alg::bbox_t> boxes;
        std::vector<object_detect_alg::bbox_t> track_boxes;
        int res = yolo->detect(&frame, boxes);
        if(res != 0) {
            std::cerr << "Error: Detection failed." << std::endl;            
        }

        DETECTIONS detections;
        for (const auto& box : boxes) {
            DETECTIONS dets = get_detections(box);
            detections.insert(detections.end(), dets.begin(), dets.end());
        }
        if(FeatureTensor::getInstance()->getRectsFeature(frame, detections)) {
            mytrack.predict();
            mytrack.update(detections);
            for(Track & track: mytrack.tracks) {
                if(!track.is_confirmed() || track.time_since_update > 1) continue;{
                    DETECTBOX box = track.to_tlwh();
                    object_detect_alg::bbox_t bbox;
                    bbox.rect = cv::Rect(box(0), box(1), box(2), box(3));
                    bbox.prob = 1.0; // 假设置信度为1.0
                    bbox.track_id = std::to_string(track.track_id);
                    track_boxes.push_back(bbox);
                }
            }
            // 比较IOU并更新检测框
            compare_iou(track_boxes, boxes);
        } else {
            std::cerr << "Error: Feature extraction failed." << std::endl;          
        }        

        // 在视频帧上绘制检测结果
        for (const auto& box : boxes) {
            cv::rectangle(frame, cv::Point(box.rect.x, box.rect.y), cv::Point(box.rect.x + box.rect.width, box.rect.y + box.rect.height), cv::Scalar(0, 255, 0), 2);
            std::string label = std::to_string(box.obj_id) + ": " + box.track_id;
            cv::putText(frame, label, cv::Point(box.rect.x, box.rect.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 2);
        }

       

        cv::imshow("Video Frame", frame);
        if (cv::waitKey(30) >= 0) break; // 等待30毫秒或按键
    }
    
    // 释放视频捕获对象
    cap.release();
    cv::destroyAllWindows();
    delete yolo;
    
    // // 初始化特征提取器
    // FeatureTensor* extractor = FeatureTensor::getInstance("./resnet50.onnx");
    
    // // 创建检测框
    // DETECTBOX box(100, 100, 100, 200); // 假设的框
    // DETECTION_ROW detectionRow;
    // detectionRow.tlwh = box;
    // detectionRow.confidence = 0.9; // 假设的置信度
    // std::vector<DETECTION_ROW> boxes;
    // boxes.push_back(detectionRow);
    
    // // 提取特征
    // extractor->getRectsFeature(testImage,boxes);
    
    // // 检查特征维度和值
    // std::cout << "特征维度: " << boxes[0].feature.size() << std::endl;
    // // std::cout << "特征范数: " << features.row(0).norm() << std::endl;
    
    // // 查看前几个特征值
    // for (int i = 0; i < 10; i++) {
    //     std::cout << boxes[0].feature[i] << " ";
    // }
    // std::cout << std::endl;
    
    return 0;
}